{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "pygame 2.1.2 (SDL 2.0.18, Python 3.8.13)\n",
      "Hello from the pygame community. https://www.pygame.org/contribute.html\n"
     ]
    },
    {
     "ename": "ModuleNotFoundError",
     "evalue": "No module named 'Envs'",
     "output_type": "error",
     "traceback": [
      "\u001B[1;31m---------------------------------------------------------------------------\u001B[0m",
      "\u001B[1;31mModuleNotFoundError\u001B[0m                       Traceback (most recent call last)",
      "Input \u001B[1;32mIn [1]\u001B[0m, in \u001B[0;36m<cell line: 4>\u001B[1;34m()\u001B[0m\n\u001B[0;32m      2\u001B[0m \u001B[38;5;28;01mimport\u001B[39;00m \u001B[38;5;21;01mpygame\u001B[39;00m\n\u001B[0;32m      3\u001B[0m \u001B[38;5;28;01mimport\u001B[39;00m \u001B[38;5;21;01mnumpy\u001B[39;00m \u001B[38;5;28;01mas\u001B[39;00m \u001B[38;5;21;01mnp\u001B[39;00m\n\u001B[1;32m----> 4\u001B[0m \u001B[38;5;28;01mfrom\u001B[39;00m \u001B[38;5;21;01mEnvs\u001B[39;00m\u001B[38;5;21;01m.\u001B[39;00m\u001B[38;5;21;01mMazeEnv\u001B[39;00m \u001B[38;5;28;01mimport\u001B[39;00m MazeEnv\n",
      "\u001B[1;31mModuleNotFoundError\u001B[0m: No module named 'Envs'"
     ]
    }
   ],
   "source": [
    "import time\n",
    "import pygame\n",
    "import numpy as np\n",
    "from Envs.MazeEnv import MazeEnv"
   ],
   "metadata": {
    "collapsed": false
   }
  },
  {
   "cell_type": "markdown",
   "source": [
    "### Agent\n",
    "* `Agent`是和环境`environment`交互的主体。\n",
    "* `predict()`方法：输入观察值`observation`（或者说状态`state`），输出动作值\n",
    "* `sample()`方法：再`predict()`方法基础上使用`ε-greedy`增加探索\n",
    "* `learn()`方法：输入训练数据，完成一轮`Q`表格的更新"
   ],
   "metadata": {
    "collapsed": false
   }
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [
    "class SarsaAgent(object):\n",
    "    def __init__(self, obs_n, act_n, learning_rate=0.01, gamma=0.9, e_greed=0.1):\n",
    "        self.act_n = act_n  # 动作维度，有几个动作可选\n",
    "        self.lr = learning_rate  # 学习率\n",
    "        self.gamma = gamma  # reward的衰减率\n",
    "        self.epsilon = e_greed  # 按一定概率随机选动作\n",
    "        self.Q = np.zeros((obs_n, act_n))\n",
    "\n",
    "    # 根据输入观察值，采样输出的动作值，带探索\n",
    "    def sample(self, obs):\n",
    "        if np.random.uniform(0, 1) < (1.0 - self.epsilon):  #根据table的Q值选动作\n",
    "            action = self.predict(obs)\n",
    "        else:\n",
    "            action = np.random.choice(self.act_n)  #有一定概率随机探索选取一个动作\n",
    "        return action\n",
    "\n",
    "    # 根据输入观察值，预测输出的动作值\n",
    "    def predict(self, obs):\n",
    "        Q_list = self.Q[obs, :]\n",
    "        maxQ = np.max(Q_list)\n",
    "        action_list = np.where(Q_list == maxQ)[0]  # maxQ可能对应多个action\n",
    "        action = np.random.choice(action_list)\n",
    "        return action\n",
    "\n",
    "    # 学习方法，也就是更新Q-table的方法\n",
    "    def learn(self, obs, action, reward, next_obs, done):\n",
    "        \"\"\" on-policy\n",
    "            obs: 交互前的obs, s_t\n",
    "            action: 本次交互选择的action, a_t\n",
    "            reward: 本次动作获得的奖励r\n",
    "            next_obs: 本次交互后的obs, s_t+1\n",
    "            next_action: 根据当前Q表格, 针对next_obs会选择的动作, a_t+1\n",
    "            done: episode是否结束\n",
    "            Q(S<t>,A<t>)=Q(S<t>,A<t>) + α[R<t+1>+γmax(Q(S<t+1>,a)) - Q(S<t>,A<t>)]\n",
    "            Q(S<t+1>,a) == 在t+1时刻,状态为S<t+1>, 所有行动的Q值\n",
    "            R<t+1> == 在t时刻采取行动之后的奖励\n",
    "            α == 学习率\n",
    "            γ == 衰减率\n",
    "        \"\"\"\n",
    "        predict_Q = self.Q[obs, action]\n",
    "        if done:\n",
    "            target_Q = reward  # 没有下一个状态了\n",
    "        else:\n",
    "            # target_Q = reward + self.gamma * self.Q[next_obs, next_action]  # Sarsa 对应 R<t+1>+γQ(S<t+1>,A<t+1>)\n",
    "            target_Q = reward + self.gamma * np.max(self.Q[next_obs, :])  # Q-learning 对应 R<t+1>+γmax(Q(S<t+1>,a))\n",
    "        self.Q[obs, action] += self.lr * (target_Q - predict_Q)  # 修正q\n",
    "\n",
    "    # 保存Q表格数据到文件\n",
    "    def save(self):\n",
    "        npy_file = './maze_game_table.npy'\n",
    "        np.save(npy_file, self.Q)\n",
    "        print(npy_file + ' saved.')\n",
    "\n",
    "    # 从文件中读取Q值到Q表格中\n",
    "    def restore(self, npy_file='./maze_game_table.npy'):\n",
    "        self.Q = np.load(npy_file)\n",
    "        print(npy_file + ' loaded.')"
   ],
   "metadata": {
    "collapsed": false
   }
  },
  {
   "cell_type": "markdown",
   "source": [
    "###  Training && Test（训练&&测试）\n",
    "* `run_episode()`：`agent`在一个`episode`中训练的过程，使用`agent.sample()`与环境交互，使用`agent.learn()`训练`Q`表格。\n",
    "* `test_episode()`：`agent`在一个`episode`中测试效果，评估目前的`agent`能在一个`episode`中拿到多少总`reward`。"
   ],
   "metadata": {
    "collapsed": false
   }
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [
    "class Training:\n",
    "    @staticmethod\n",
    "    def run_episode(env, agent, render=False):\n",
    "        total_steps = 0  # 记录每个episode走了多少step\n",
    "        total_reward = 0\n",
    "\n",
    "        obs = env.reset()  # 重置环境, 重新开一局（即开始新的一个episode）\n",
    "        action = agent.sample(obs)  # 根据算法选择一个动作\n",
    "\n",
    "        while True:\n",
    "            for event in pygame.event.get():\n",
    "                if event.type == pygame.QUIT:\n",
    "                    pygame.quit()\n",
    "            next_obs, reward, done, _ = env.step(action, render)  # 与环境进行一个交互\n",
    "            next_action = agent.sample(next_obs)  # 根据算法选择一个动作\n",
    "            # 训练 Sarsa 算法\n",
    "            agent.learn(obs, action, reward, next_obs, done)\n",
    "\n",
    "            action = next_action\n",
    "            obs = next_obs  # 存储上一个观察值\n",
    "            total_reward += reward\n",
    "            total_steps += 1  # 计算step数\n",
    "            if render:\n",
    "                # time.sleep(0.01)\n",
    "                env.gameMaze.draw_maze(env.maze, env.cur_pos)\n",
    "            if done:\n",
    "                break\n",
    "        return total_reward, total_steps\n",
    "\n",
    "    @staticmethod\n",
    "    def test_episode(env, agent, render=False):\n",
    "        total_reward = 0\n",
    "        obs = env.reset()\n",
    "        total_step = 0\n",
    "        while True:\n",
    "            for event in pygame.event.get():\n",
    "                if event.type == pygame.QUIT:\n",
    "                    pygame.quit()\n",
    "            action = agent.predict(obs)  # greedy\n",
    "            next_obs, reward, done, _ = env.step(action, render)\n",
    "            total_reward += reward\n",
    "            obs = next_obs\n",
    "            if render:\n",
    "                time.sleep(0.05)\n",
    "                env.gameMaze.draw_maze(env.maze, env.cur_pos)\n",
    "            # env.render()\n",
    "            if done:\n",
    "                break\n",
    "            total_step += 1\n",
    "            if total_step > (env.size // 2) ** 2:\n",
    "                break\n",
    "        return total_reward\n"
   ],
   "metadata": {
    "collapsed": false
   }
  },
  {
   "cell_type": "markdown",
   "source": [
    "### 创建环境和Agent，启动训练"
   ],
   "metadata": {
    "collapsed": false
   }
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [
    "env = MazeEnv(\n",
    "    minsize=8,\n",
    "    maxsize=8,\n",
    ")\n",
    "\n",
    "agent = SarsaAgent(\n",
    "    obs_n=env.observation_space.n,\n",
    "    act_n=env.action_space.n,\n",
    "    learning_rate=0.1,\n",
    "    gamma=0.9,\n",
    "    e_greed=0.1\n",
    ")"
   ],
   "metadata": {
    "collapsed": false
   }
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [
    "env.gameMaze.start_game()\n",
    "time.sleep(2)"
   ],
   "metadata": {
    "collapsed": false
   }
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [
    "# 训练3000个episode，打印每个episode的分数\n",
    "# agent.restore()\n",
    "start_time = time.time()\n",
    "\n",
    "for episode in range(4000):\n",
    "    # render = False\n",
    "    # if episode !=0 and episode % 50 == 0:\n",
    "    #     render=True\n",
    "    render = True\n",
    "    ep_reward, ep_steps = Training.run_episode(env, agent, render)\n",
    "    print(f'\\r#######Episode {episode}: steps={ep_steps}, reward={ep_reward}', end='#######')\n",
    "end_time = time.time()\n",
    "print()\n",
    "print(f'cost time:{end_time - start_time}')\n",
    "agent.save()"
   ],
   "metadata": {
    "collapsed": false
   }
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [
    "# 全部训练结束，查看算法效果\n",
    "# agent.restore()\n",
    "time.sleep(5)\n",
    "test_reward = Training.test_episode(env, agent, True)\n",
    "print('test reward = %.1f' % (test_reward))\n"
   ],
   "metadata": {
    "collapsed": false
   }
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [
    "pygame.quit()"
   ],
   "metadata": {
    "collapsed": false
   }
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 2
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython2",
   "version": "2.7.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 0
}
